/**
  ******************************************************************************
  * @file    iap.c
  * @author  Puya Application Team
  * @brief   iap application entry point
  ******************************************************************************
  * @attention
  *
  * <h2><center>&copy; Copyright (c) 2021 Puya Semiconductor.
  * All rights reserved.</center></h2>
  *
  *
  ******************************************************************************
  */


/* Includes ------------------------------------------------------------------*/
#include "main.h"
#include "iap.h"
#include "usart.h"
#include "flash.h"
#include "wdg.h"
#include "common.h"

#pragma pack(1)
typedef enum
{
  AREA_ERROR = 0x00U, /* Error Address Area */
  FLASH_AREA = 0x01U, /* Flash Address Area */
  RAM_AREA   = 0x02U, /* RAM Address area */
  OB_AREA    = 0x03U, /* Option bytes Address area */
  OTP_AREA   = 0x04U, /* OTP Address area */
  SYS_AREA   = 0x05U, /* System memory area */
  EB_AREA    = 0x07U  /* Engi bytes Address area */
} AreaType_T;
#pragma pack()

const uint8_t BOOTLOADER_CMD[] =
{
  CMD_GET_COMMAND,
  CMD_GET_VERSION,
  CMD_GET_ID,
  CMD_SET_BR,
  CMD_READ_MEMORY,
  CMD_GO,
  CMD_WRITE_MEMORY,
  CMD_EXT_ERASE_MEMORY,
  CMD_WRITE_PROTECT,
  CMD_WRITE_UNPROTECT,
  CMD_READ_PROTECT,
  CMD_READ_UNPROTECT,
  CMD_GET_DEVICE_IDCODE,
};


/* Private variables ---------------------------------------------------------*/
uint8_t gBuffer[256 + 1] = { 0 };

typedef void (*Function_Pointer)(void);

/**
 * @brief  Get address area type.
 * @param  puiAddr  Address to check
 * @retval FLASH_AREA, RAM_AREA, ... or AREA_ERROR
 */
static uint8_t GetAddressArea(uint32_t *puiAddr)
{
  uint8_t ubData[4] = { 0U };

  ubData[0] = UsartReadByte();
  ubData[1] = UsartReadByte();
  ubData[2] = UsartReadByte();
  ubData[3] = UsartReadByte();

  if (UsartReadByte() != GetXOR(ubData, 0x04, 0x00))
  {
    return AREA_ERROR;
  }

  UsartSendByte(ACK_BYTE);
  *puiAddr = (ubData[0] << 24) + (ubData[1] << 16) + \
             (ubData[2] << 8) + ubData[3];

  uint32_t uiStartAddr = FLASH_BASE;
  uint32_t uiEndAddr   = FLASH_BASE + 0x10000;

  if ((*puiAddr >= uiStartAddr) && (*puiAddr < uiEndAddr))
  {
    return FLASH_AREA;
  }

  uiStartAddr = SRAM_BASE + 0x0200;
  uiEndAddr   = SRAM_BASE + 0x02000;

  if ((*puiAddr >= uiStartAddr) && (*puiAddr < uiEndAddr))
  {
    return RAM_AREA;
  }

  uiStartAddr = SYSTEM_BASE;
  uiEndAddr   = SYSTEM_BASE + 0x0E00;

  if ((*puiAddr >= uiStartAddr) && (*puiAddr < uiEndAddr))
  {
    return SYS_AREA;
  }

  uiStartAddr = SYSTEM_BASE + 0x0E00;
  uiEndAddr   = 0x1FFF1180;

  if ((*puiAddr >= uiStartAddr) && (*puiAddr < uiEndAddr))
  {
    return OB_AREA;
  }

  return AREA_ERROR;
}

/**
 * @brief  Read memory from device.
 * @retval SUCCESS or ERROR
 */
static ErrorStatus ReadMemory(void)
{
  uint32_t uiAddr = 0UL;

  if (AREA_ERROR == GetAddressArea(&uiAddr))
  {
    return ERROR;
  }

  uint8_t ubData = UsartReadByte();

  if ((ubData ^ UsartReadByte()) != 0xFFU)
  {
    return ERROR;
  }

  UsartSendByte(ACK_BYTE);

  for (uint16_t i = 0U; i < ubData + 1; i++)
  {
    UsartSendByte(M8(uiAddr + i));
  }

  return SUCCESS;
}

/**
 * @brief  Jump to user application.
 * @retval SUCCESS or ERROR
 */
static ErrorStatus JumpToAddr(void)
{
  uint32_t uiAddr = 0U;
  uint8_t ubType = GetAddressArea(&uiAddr);

  if ((ubType != FLASH_AREA) && (ubType != RAM_AREA))
  {
    return ERROR;
  }

  WRITE_REG(RCC->IOPENR,   0x00000000UL);
  WRITE_REG(RCC->APBENR2,  0x00000000UL);
  WRITE_REG(USART1->CR1,   0x00000000UL);
  WRITE_REG(USART1->BRR,   0x00000000UL);
  WRITE_REG(GPIOA->MODER,  0xEBFFFFFFUL);
  WRITE_REG(GPIOA->AFR[0], 0x00000000UL);
  WRITE_REG(GPIOA->AFR[1], 0x00000000UL);

  CLEAR_BIT(RCC->CFGR, RCC_CFGR_SW);

  while (0 != READ_BIT(RCC->CFGR, RCC_CFGR_SWS))
  {
    WdgRefresh();
  }

  CLEAR_BIT(RCC->CR, RCC_CR_PLLON);

  __enable_irq();

  /* Set Main Stack Pointer */
  __set_MSP(M32(uiAddr));

  ((Function_Pointer)M32(uiAddr + 4U))();

  return SUCCESS;
}

/**
 * @brief  Write data to device memory.
 * @retval SUCCESS or ERROR
 */
static ErrorStatus WriteMemory(void)
{
  uint32_t uiAddr = 0UL;
  ErrorStatus enFlag = SUCCESS;

  uint8_t ubType = GetAddressArea(&uiAddr);

  if ((ubType != FLASH_AREA) && \
      (ubType != RAM_AREA)   && \
      (ubType != OB_AREA))
  {
    return ERROR;
  }

  uint8_t ubLength = UsartReadByte();

  for (uint16_t i = 0; i < ubLength + 1; i++)
  {
    gBuffer[i] = UsartReadByte();
  }

  if (UsartReadByte() != GetXOR(gBuffer, ubLength + 1, ubLength))
  {
    return ERROR;
  }

  switch (ubType)
  {
    case FLASH_AREA:
      enFlag = WriteFlash(uiAddr, gBuffer, ubLength);
      break;

    case RAM_AREA:
      for (uint16_t i = 0; i < ubLength + 1; i++)
      {
        M8(uiAddr + i) = gBuffer[i];
      }

      enFlag = SUCCESS;
      break;

    case OB_AREA:
      enFlag = WriteOption(uiAddr, gBuffer, ubLength);
      break;
  }

  return enFlag;
}

/**
 * @brief  Erase memory.
 * @retval SUCCESS or ERROR
 */
static ErrorStatus ExtendedErase(void)
{
  uint8_t ubDataTemp = 0U;
  ErrorStatus enFlag = ERROR;

  /* Read number of pages to be erased */
  uint8_t ubEraseType = UsartReadByte();
  uint8_t ubLength = UsartReadByte();

  /* Checksum initialization */
  uint8_t ubXorVal  = ubEraseType;
  ubXorVal ^= ubLength;

  if (0xFFU == ubEraseType)//0xFFFY, Y=F,E,D
  {
    /* eceive double-byte checksum */
    if (UsartReadByte() != ubXorVal)
    {
      return ERROR;
    }

    return MassErase();//0xFFFF = Mass erase
  }
  else
  {
    for (uint8_t i = 0; i < 2 * (ubLength + 1); i++)
    {
      gBuffer[i] = UsartReadByte();
    }

    ubXorVal = GetXOR(gBuffer, 2 * (ubLength + 1), ubXorVal);

    if (UsartReadByte() != ubXorVal)
    {
      return ERROR;
    }

    for (uint8_t i = 0; i < ubLength + 1; i++)
    {
      ubDataTemp = gBuffer[2 * i + 0];
      gBuffer[2 * i + 0] = gBuffer[2 * i + 1];
      gBuffer[2 * i + 1] = ubDataTemp;
    }

    switch (ubEraseType)
    {
      case 0x00://Compatible with ST, erase 1KBytes at a time
        enFlag = PageErase((uint16_t *)gBuffer, ubLength, 8);
        break;

      case 0x10://PageErase
        enFlag = PageErase((uint16_t *)gBuffer, ubLength, 1);
        break;

      case 0x20://SectorErase
        enFlag = SectorErase((uint16_t *)gBuffer, ubLength);
        break;

      default:
        enFlag = ERROR;
        break;
    }
  }

  return enFlag;
}

/**
 * @brief  Enable write protection.
 * @retval SUCCESS or ERROR
 */
static ErrorStatus WriteProtect(void)
{
  uint16_t ubProtectedPages = 0xFFFF;

  uint8_t ubLength = UsartReadByte();

  for (uint8_t i = 0; i < ubLength + 1; i++)
  {
    gBuffer[i] = UsartReadByte();
  }

  if (UsartReadByte() != GetXOR(gBuffer, ubLength + 1, ubLength))
  {
    return ERROR;
  }

  for (uint8_t i = 0; i < ubLength + 1; i++)
  {
    CLEAR_BIT(ubProtectedPages, (1U << gBuffer[i]));
  }

  gBuffer[0] = ubProtectedPages & 0xFFU;
  gBuffer[1] = (ubProtectedPages >> 8U) & 0xFFU;

  return WriteOption(WRPR_BASE, gBuffer, 0x01);
}

/**
 * @brief  Check if command is valid.
 * @param  Command code
 */
static void IsValidCommand(uint8_t ubCmd)
{
  switch (ubCmd)
  {
    case CMD_WRITE_MEMORY:
    case CMD_EXT_ERASE_MEMORY:
    case CMD_READ_MEMORY:
    case CMD_GO:
    case CMD_WRITE_PROTECT:
    case CMD_WRITE_UNPROTECT:
    case CMD_READ_PROTECT:

    case CMD_GET_COMMAND:
    case CMD_GET_VERSION:
    case CMD_GET_ID:
    case CMD_GET_DEVICE_IDCODE:
    case CMD_READ_UNPROTECT:
    case CMD_SET_BR:
      UsartSendByte(ACK_BYTE);
      break;
  }
}

/**
 * @brief  Unlock FLASH for operation.
 * @param  Command code
 * @retval SUCCESS or ERROR
 */
static ErrorStatus FlashUnlock(uint8_t ubCmd)
{
  ErrorStatus enStatus = SUCCESS;

  /* FLASH_Unlock (Unified entry, prevents the chip from running wild,
  *  avoids incorrect FLASH operations)
  */
  switch (ubCmd)
  {
    case CMD_WRITE_MEMORY:
    case CMD_EXT_ERASE_MEMORY:
    case CMD_WRITE_PROTECT:
    case CMD_WRITE_UNPROTECT:
    case CMD_READ_PROTECT:
    case CMD_READ_UNPROTECT:
      if (FLASH_SR_BSY == READ_BIT(FLASH->SR, FLASH_SR_BSY))
      {
        UsartSendByte(NACK_BYTE);

        enStatus = ERROR;

        break;
      }

      if (FLASH_CR_LOCK == READ_BIT(FLASH->CR, FLASH_CR_LOCK))
      {
        WRITE_REG(FLASH->KEYR, FLASH_KEY1);
        WRITE_REG(FLASH->KEYR, FLASH_KEY2);
      }
  }

  return enStatus;
}

/**
 * @brief  Lock FLASH.
 * @param  Command code
 */
static void FlashLock(uint8_t ubCmd)
{
  switch (ubCmd)
  {
    case CMD_WRITE_MEMORY:
    case CMD_EXT_ERASE_MEMORY:
    case CMD_WRITE_PROTECT:
    case CMD_WRITE_UNPROTECT:
    case CMD_READ_PROTECT:
    case CMD_READ_UNPROTECT:
      SET_BIT(FLASH->CR, FLASH_CR_LOCK);
      break;
  }
}

/**
 * @brief  Execute the specified command.
 * @param  Command code
 * @retval SUCCESS or ERROR
 */
static ErrorStatus ExecuteCommand(uint8_t ubCmd)
{
  ErrorStatus enStatus = SUCCESS;
  uint8_t i = 0;

  switch (ubCmd)
  {
    case CMD_GET_COMMAND:
      UsartSendByte(COUNTOF(BOOTLOADER_CMD));
      UsartSendByte(VERSION);

      for (i = 0; i < COUNTOF(BOOTLOADER_CMD); i++)
      {
        UsartSendByte(BOOTLOADER_CMD[i]);
      }

      break;

    case CMD_GET_VERSION:
      UsartSendByte(VERSION);
      UsartSendByte(0x00);
      UsartSendByte(0x00);
      break;

    case CMD_GET_ID:
      UsartSendByte(0x01);
      UsartSendByte(0x01);//0106
      UsartSendByte(0x06);
      break;

    case CMD_GET_DEVICE_IDCODE:
      UsartSendByte(0x01);
      UsartSendByte(DBGMCU_IDCODE >> 8);
      UsartSendByte(DBGMCU_IDCODE & 0xFF);
      break;

    case CMD_GO:
      enStatus = JumpToAddr();
      break;

    case CMD_WRITE_MEMORY:
      enStatus = WriteMemory();
      break;

    case CMD_EXT_ERASE_MEMORY:
      enStatus = ExtendedErase();
      break;

    case CMD_WRITE_PROTECT:
      enStatus = WriteProtect();
      break;

    case CMD_WRITE_UNPROTECT:
      gBuffer[0] = (FLASH_WRPR >> 0) & 0xFF;
      gBuffer[1] = (FLASH_WRPR >> 8) & 0xFF;
      enStatus = WriteOption(WRPR_BASE, gBuffer, 0x01);
      break;

    case CMD_READ_PROTECT:
      gBuffer[0] = FLASH_OPTR_RDP_LEVEL_1;
      enStatus = WriteOption(OPTR1_BASE, gBuffer, 0x00);
      break;

    case CMD_READ_UNPROTECT:
      gBuffer[0] = FLASH_OPTR_RDP_LEVEL_0;
      enStatus = WriteOption(OPTR1_BASE, gBuffer, 0x00);
      break;

    default:
      enStatus = ERROR;
      break;
  }

  return enStatus;
}

/**
 * @brief  Initialize Bootloader.
 */
void IAPInit(void)
{
  WdgInit();

  UsartInit();

  while(1)
  {
    WdgRefresh();

    if (UsartReadByte() == START_BYTE)
    {
      break;
    }
  }

  CLEAR_BIT(USART1->SR, (USART_SR_RXNE | USART_SR_TC));

  UsartSendByte(ACK_BYTE);
}

/**
 * @brief  IAP process
 */
void IAPProcess(void)
{
  ErrorStatus enStatus = SUCCESS;

  uint8_t ubCmd = UsartReadByte();

  /* Check the data integrity */
  if ((ubCmd ^ UsartReadByte()) != 0xFF)
  {
    UsartSendByte(NACK_BYTE);
    return;
  }

  IsValidCommand(ubCmd);

  /* FLASH_Unlock (Unified entry, prevents the chip
  *  from running wild, avoids incorrect FLASH operations)
  */
  if(ERROR == FlashUnlock(ubCmd))
  {
    return ;
  }

  if(ubCmd == CMD_READ_MEMORY)
  {
    (void)ReadMemory();
    return;
  }

  if(ubCmd == CMD_SET_BR)
  {
    UsartConfigureBps();
    return ;
  }

  enStatus = ExecuteCommand(ubCmd);

  FlashLock(ubCmd);

  UsartSendByte((SUCCESS == enStatus) ? ACK_BYTE : NACK_BYTE);
}

/**
 * @brief  Jump to the specified address and run
 */
void JumpToAddress(uint32_t uiAddr)
{
  uint32_t uiSP = M32(uiAddr);
  uint32_t uiPC = M32(uiAddr + 4);

  if ((SRAM_BASE == (uiSP & 0xFF000000UL)) &&\
      ((uiPC & 0xFF000000UL) == FLASH_BASE))
  {
    __enable_irq();
    
    /* Set Main Stack Pointer */
    __set_MSP(M32(uiAddr));

    ((Function_Pointer)M32(uiAddr + 4U))();
  }
}
